import os
import requests
import baostock as bs  # 需要 pip install baostock

# 下载全部上市公司的股票图标
# 下载目录
DOWNLOAD_DIR = "static/images"

# 确保下载目录存在
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

def download_stock_images():
    """使用 baostock 获取沪深 A 股列表，并下载股票图标"""
    try:
        # baostock 登录
        lg = bs.login()
        if lg.error_code != '0':
            raise Exception(f"baostock login failed: {lg.error_msg}")
        
        # 获取所有股票
        rs = bs.query_all_stock(day="2025-10-22")  # 使用当前日期
        if rs.error_code != '0':
            raise Exception(f"query_all_stock failed: {rs.error_msg}")
        
        # 提取股票代码
        stock_list = []
        while rs.next():
            stock_list.append(rs.get_row_data())
        
        # 过滤沪市 (sh.) 和深市 (sz.)
        sh_codes = [row[0].split('.')[1] for row in stock_list if row[0].startswith('sh.')]
        sz_codes = [row[0].split('.')[1] for row in stock_list if row[0].startswith('sz.')]
        
        # 额外过滤基于代码前缀（避免错误分配）
        sh_stocks = [code for code in sh_codes if code.startswith(('60', '68'))]
        sz_stocks = [code for code in sz_codes if code.startswith(('00', '30', '002'))]
        
        # 调试：打印前10个
        print("调试: 前10个沪市代码:", sh_stocks[:10])
        print("调试: 前10个深市代码:", sz_stocks[:10])
        
        print(f"🎯 找到 {len(sh_stocks)} 个沪市股票和 {len(sz_stocks)} 个深市股票，开始下载图标...")
        
        success_count = 0
        total = len(sh_stocks) + len(sz_stocks)
        
        # 下载沪市图标
        for i, code in enumerate(sh_stocks, 1):
            url = f"https://sns.sseinfo.com/resources/images/avatar/company/{code}.png"
            filename = f"{code}.png"
            if download_image(url, filename):
                success_count += 1
            print_progress(i, total, code, "沪市")
        
        # 下载深市图标
        for i, code in enumerate(sz_stocks, len(sh_stocks) + 1):
            url = f"https://resstatic.cninfo.com.cn/irm/ssgs/S{code}/images/{code}.gif"
            filename = f"{code}.gif"
            if download_image(url, filename):
                success_count += 1
            print_progress(i, total, code, "深市")
        
        print(f"\n🎉 下载完成！成功 {success_count}/{total} 个文件")
        print(f"📁 保存目录: {DOWNLOAD_DIR}")
        
    except Exception as e:
        print(f"❌ 获取股票列表失败: {e}")

def download_image(url: str, filename: str) -> bool:
    """下载单个图像文件"""
    try:
        filepath = os.path.join(DOWNLOAD_DIR, filename)
        response = requests.get(url, timeout=10)
        if response.status_code == 404:
            print(f"⚠️ 404 未找到: {filename} (跳过)")
            return False
        response.raise_for_status()
        with open(filepath, 'wb') as f:
            f.write(response.content)
        return True
    except Exception as e:
        print(f"❌ 下载失败 {filename}: {e}")
        return False

def print_progress(current: int, total: int, code: str, market: str):
    """打印下载进度"""
    percent = (current / total) * 100
    print(f"[{current}/{total}] ({percent:.1f}%) {market} {code} 下载中...")

if __name__ == "__main__":
    download_stock_images()